package com.dddframework.web.utils;

import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NoArgsConstructor;
import lombok.SneakyThrows;
import lombok.experimental.UtilityClass;
import lombok.extern.slf4j.Slf4j;

import javax.websocket.*;
import javax.websocket.server.ServerContainer;
import java.io.IOException;
import java.net.URI;
import java.net.URISyntaxException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Consumer;

/**
 * WebSocket工具
 */
@Slf4j(topic = "### BASE-WEB : WebSocketKit ###")
@UtilityClass
public class WebSocketKit {
    private static final Map<String, Session> SESSION_POOL = new ConcurrentHashMap<>();

    /**
     * 启动WebSocket服务端
     */
    public void startServer(Class<? extends BaseWebSocketServer> serverClass) {
        try {
            ServerContainer container = (ServerContainer) ContainerProvider.getWebSocketContainer();
            container.addEndpoint(serverClass);
            log.info("WebSocket server started [{}]", serverClass.getSimpleName());
        } catch (DeploymentException e) {
            log.error("Failed to start WebSocket server [{}]", serverClass.getSimpleName(), e);
        }
    }

    /**
     * 初始化WebSocket客户端并建立连接
     */
    public <T> void connectServer(String biz, String url, Consumer<T> listener, Consumer<Session> onOpen) {
        try {
            WebSocketContainer container = ContainerProvider.getWebSocketContainer();
            Session session = container.connectToServer(new WebSocketClientEndpoint<>(listener, onOpen), new URI(url));
            if (!SESSION_POOL.containsKey(biz)) {
                SESSION_POOL.put(biz, session);
            } else {
                log.error("WebSocket already listening: {}", url);
            }
            log.info("WebSocket connection established for biz: {}", biz);
        } catch (DeploymentException | IOException e) {
            throw new RuntimeException("Failed to establish WebSocket connection.", e);
        } catch (URISyntaxException e) {
            throw new RuntimeException("Invalid WebSocket URL: " + url, e);
        }
    }

    public <T> void connectServer(String url, Consumer<T> listener, Consumer<Session> onOpen) {
        connectServer("DEFAULT", url, listener, onOpen);
    }

    public <T> void connectServer(String url, Consumer<T> listener) {
        connectServer("DEFAULT", url, listener, null);
    }

    /**
     * 发送文本消息到WebSocket服务器。
     *
     * @param biz  业务域
     * @param text 要发送的消息
     */
    public void sendText(String biz, String text) {
        Session session = SESSION_POOL.get(biz);
        if (session != null && session.isOpen()) {
            session.getAsyncRemote().sendText(text);
            log.info("Message sent to WebSocket server for biz {}: {}", biz, text);
        } else {
            log.warn("WebSocket session is not open for biz: {}", biz);
        }
    }

    /**
     * 发送文本消息到WebSocket服务器。
     */
    public void sendText(String text) {
        sendText("DEFAULT", text);
    }

    /**
     * 发送对象消息到WebSocket服务器。
     *
     * @param biz    业务域
     * @param object 要发送的消息
     */
    public void sendObject(String biz, Object object) {
        Session session = SESSION_POOL.get(biz);
        if (session != null && session.isOpen()) {
            session.getAsyncRemote().sendObject(object);
            log.info("Message sent to WebSocket server for biz {}: {}", biz, object);
        } else {
            log.warn("WebSocket session is not open for biz: {}", biz);
        }
    }

    public void sendObject(Object object) {
        sendObject("DEFAULT", object);
    }

    /**
     * 断开连接
     */
    @SneakyThrows
    public void disconnect(String biz) {
        Session session = SESSION_POOL.get(biz);
        if (session != null) {
            session.close();
        }
    }

    /**
     * WebSocket endpoint implementation.
     */
    @Data
    @AllArgsConstructor
    @NoArgsConstructor
    @ClientEndpoint
    public static class WebSocketClientEndpoint<T> implements MessageHandler.Whole<T> {
        private Consumer<T> onMessage;
        private Consumer<Session> onOpen;

        @Override
        public void onMessage(T message) {
            log.info("Received ws message: {}", message);
            onMessage.accept(message);
        }

        @OnOpen
        public void onOpen(Session session) {
            log.info("WebSocket connection opened for session: {}", session.getId());
            if (onOpen != null) {
                onOpen.accept(session);
            }
        }

        @OnClose
        public void onClose(Session session, CloseReason reason) {
            log.info("WebSocket connection closed for session: {}, reason: {}", session.getId(), reason);
            SESSION_POOL.remove(session.getId());
        }

        @OnError
        public void onError(Throwable throwable) {
            log.error("WebSocket error", throwable);
        }
    }
}